# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
容器发射器
处理列表、字典、集合、元组等容器类型的代码生成。
"""

from typing import Callable, Dict, Optional, Type, Union
from ...ir import IRNode
from ...ir.expressions import IRList, IRDict, IRSet, IRTuple
from ..core.emitter import NodeTypeEmitter, EmitResult
from ..core.context import CodeGenContext
from ..core.generator import get_global_generator


class ContainerEmitter(NodeTypeEmitter):
    """
    容器发射器

    处理IRList, IRDict, IRSet, IRTuple节点的代码生成。
    """

    def __init__(self):
        super().__init__([IRList, IRDict, IRSet, IRTuple])
        # 使用类型分发映射，避免多个if-elif判断
        self._emitter_map: Dict[
            Type[IRNode], Callable[[IRNode, CodeGenContext], EmitResult]
        ] = {
            IRList: self._emit_list,
            IRDict: self._emit_dict,
            IRSet: self._emit_set,
            IRTuple: self._emit_tuple,
        }

    def emit(self, node: IRNode, context: CodeGenContext) -> EmitResult:
        """生成容器代码"""
        # 使用类型分发映射
        for node_type, emitter_func in self._emitter_map.items():
            if isinstance(node, node_type):
                return emitter_func(node, context)

        return EmitResult(
            code=f"/* Unknown container type: {type(node).__name__} */", precedence=0
        )

    def _get_container_base_template(self, cpp_type: Optional[str]) -> Optional[str]:
        """从 C++ 类型中提取容器的基础模板名

        例如：
        - mc::array<int> -> mc::array
        - QList<QString> -> QList
        - mc::list<T> -> mc::list

        这样就不会硬编码特定的容器类型。

        Args:
            cpp_type: 容器的 C++ 类型

        Returns:
            基础模板名，如果无法提取则返回 None
        """
        from mcpy.type_system import extract_base_template_name

        if not cpp_type:
            return None
        return extract_base_template_name(cpp_type)

    def _generate_container_code(
        self, cpp_type: str, element_codes: list, can_use_init_list: bool
    ) -> str:
        """生成容器初始化代码的通用方法

        根据是否可以使用初始化列表来决定代码格式：
        - 可用初始化列表：{e1, e2, ...} 或 {}（空容器）
        - 不可用初始化列表：cpp_type{e1, e2, ...} 或 cpp_type{}（空容器）
        - 当单行代码超过 140 字符时自动换行，提高可读性

        Args:
            cpp_type: C++ 类型（保证非空）
            element_codes: 元素代码列表
            can_use_init_list: 是否可以使用初始化列表（省略类型）

        Returns:
            生成的容器初始化代码
        """
        # 先尝试生成单行代码
        if can_use_init_list:
            single_line = f"{{{', '.join(element_codes)}}}"
        else:
            single_line = f"{cpp_type}{{{', '.join(element_codes)}}}"

        if len(single_line) <= 140:
            return single_line

        # 否则换行显示（提高可读性）
        elements_str = ",\n        ".join(element_codes)
        if can_use_init_list:
            return f"{{\n        {elements_str}\n    }}"
        else:
            return f"{cpp_type}{{\n        {elements_str}\n    }}"

    def _extract_element_type(self, cpp_type: Optional[str]) -> Optional[str]:
        """从容器类型中提取元素类型"""
        if not cpp_type:
            return None

        # 通用方法：提取 <...> 中的内容
        template_start = cpp_type.find("<")
        if template_start < 0 or not cpp_type.endswith(">"):
            return None

        # 提取 <...> 中的内容
        inner = cpp_type[template_start + 1 : -1]

        # 使用括号计数验证（处理嵌套模板）
        bracket_count = 0
        for char in inner:
            if char == "<":
                bracket_count += 1
            elif char == ">":
                bracket_count -= 1
                if bracket_count < 0:
                    # 括号不匹配
                    return None

        if bracket_count != 0:
            # 括号不匹配
            return None

        return inner

    def _emit_list(
        self, node: Union[IRNode, IRList], context: CodeGenContext
    ) -> EmitResult:
        """生成列表代码"""

        generator = get_global_generator()

        # 类型断言
        list_node = node if isinstance(node, IRList) else None
        if list_node is None:
            return EmitResult(
                code=f"/* Invalid list node type: {type(node).__name__} */",
                precedence=0,
            )

        # 获取目标类型
        cpp_type = context.get_cpp_type(node)

        # 检查是否在 dict 字面值的值位置
        from mcpy.ir.expressions import IRDict, IRCall

        is_in_dict_value = False
        is_in_constructor_call = False
        if hasattr(list_node, "parent") and list_node.parent:
            parent = list_node.parent
            if isinstance(parent, IRDict):
                # 检查当前列表是否在 dict 的 values 中
                is_in_dict_value = list_node in parent.values
            elif isinstance(parent, IRCall):
                # 检查是否是构造函数调用（如 list_construct, dict_construct）
                func_name = parent.get_func_name()
                if func_name and func_name.endswith("_construct"):
                    is_in_constructor_call = True
                else:
                    # 检查是否是类型构造函数调用（如 list[int]()）
                    # 如果 func 是 IRSubscript，说明是类型构造函数调用
                    from mcpy.ir.expressions.access import IRSubscript
                    if isinstance(parent.func, IRSubscript):
                        is_in_constructor_call = True

        # 提取元素的期望类型（用于嵌套容器）
        element_expected_type = self._extract_element_type(cpp_type)

        # 生成元素代码
        element_codes = []
        for element in list_node.elements:
            element_context = context.create_child_context(
                current_node=element,
                expected_cpp_type=element_expected_type,  # 传递元素期望类型
            )
            element_code = generator.generate_expression(element, element_context)
            element_codes.append(element_code)

        if not cpp_type:
            # 默认使用弱类型数组 mc::variants
            cpp_type = "mc::variants"

        # 获取期望类型（智能决策：何时可以使用初始化列表）
        # 1. 赋值上下文：期望类型匹配时可以省略类型
        # 2. 方法调用参数：当期望类型与实际类型完全匹配时，可以使用初始化列表
        #    （C++ 可以通过类型推导转换，如 {1, 2} -> mc::variants）
        # 3. 构造函数调用参数：必须使用显式类型（如 list_construct 的参数）
        expected_cpp_type = context.expected_cpp_type
        can_use_init_list = bool(
            expected_cpp_type
            and expected_cpp_type == cpp_type
            and not is_in_dict_value
            and not is_in_constructor_call  # 构造函数调用必须显式类型
        )

        # 生成代码
        code = self._generate_container_code(cpp_type, element_codes, can_use_init_list)
        return EmitResult(code=code, precedence=0)

    def _emit_dict(
        self, node: Union[IRNode, IRDict], context: CodeGenContext
    ) -> EmitResult:
        """生成字典代码"""

        generator = get_global_generator()

        # 类型断言
        dict_node = node if isinstance(node, IRDict) else None
        if dict_node is None:
            return EmitResult(
                code=f"/* Invalid dict node type: {type(node).__name__} */",
                precedence=0,
            )

        # 获取目标类型
        cpp_type = context.get_cpp_type(node)

        # 获取字典的值类型（从 inferred_type 中提取）
        value_expected_cpp_type = None
        inferred_type = node.get_extension_param("inferred_type")
        if inferred_type and inferred_type.name == "dict" and inferred_type.args and len(inferred_type.args) >= 2:
            # 提取值类型（第二个类型参数）
            value_type_info = inferred_type.args[1]
            # 转换为 C++ 类型
            from mcpy.type_system import get_cpp_type
            value_expected_cpp_type = get_cpp_type(value_type_info, node)

        # 生成键值对代码
        pairs = []
        for key, value in zip(dict_node.keys, dict_node.values):
            key_context = context.create_child_context(current_node=key)
            # dict 初始化列表中的值需要显式类型标记，不能省略
            # 设置期望类型，让值表达式（如元组）能够正确生成
            value_context = context.create_child_context(
                current_node=value, 
                in_assignment=False,
                expected_cpp_type=value_expected_cpp_type
            )

            key_code = generator.generate_expression(key, key_context)
            value_code = generator.generate_expression(value, value_context)
            pairs.append(f"{{{key_code}, {value_code}}}")

        if not cpp_type:
            cpp_type = "mc::dict"

        # 获取容器基础模板名
        base_template = self._get_container_base_template(cpp_type)

        # 在赋值上下文中，非空字典可以省略类型前缀（仅限模板容器）
        in_assignment = context.in_assignment

        # 只有模板容器（有类型参数）在赋值上下文中才能省略类型
        # 非模板类型（如 mc::dict）不能省略，因为 base_template == cpp_type
        is_template = base_template is not None and base_template != cpp_type
        can_use_init_list = bool(in_assignment and is_template)

        # 生成代码
        code = self._generate_container_code(cpp_type, pairs, can_use_init_list)
        return EmitResult(code=code, precedence=0)

    def _emit_set(
        self, node: Union[IRNode, IRSet], context: CodeGenContext
    ) -> EmitResult:
        """生成集合代码"""

        generator = get_global_generator()

        # 类型断言
        set_node = node if isinstance(node, IRSet) else None
        if set_node is None:
            return EmitResult(
                code=f"/* Invalid set node type: {type(node).__name__} */", precedence=0
            )

        # 获取目标类型
        cpp_type = context.get_cpp_type(node)

        if not cpp_type:
            cpp_type = "std::unordered_set<mc::variant>"

        element_expected_type = self._extract_element_type(cpp_type)

        # 生成元素代码
        element_codes = []
        for element in set_node.elements:
            element_context = context.create_child_context(
                current_node=element,
                expected_cpp_type=element_expected_type,  # 传递元素期望类型
            )
            element_code = generator.generate_expression(element, element_context)
            element_codes.append(element_code)

        # 获取容器基础模板名
        base_template = self._get_container_base_template(cpp_type)

        # 在赋值上下文中，如果期望类型与实际类型一致，可以省略类型前缀
        in_assignment = context.in_assignment
        expected_cpp_type = context.expected_cpp_type

        # 智能判断是否可以使用初始化列表：
        # 1. 必须在赋值上下文中
        # 2. 必须是模板容器
        # 3. 期望类型必须与实际类型兼容
        can_use_init_list = False
        if in_assignment and base_template is not None:
            if expected_cpp_type == cpp_type:
                # 类型完全匹配
                can_use_init_list = True
            elif expected_cpp_type and self._is_compatible_container_types(
                expected_cpp_type, cpp_type
            ):
                # 容器类型兼容（如强类型容器可以转换为弱类型容器）
                can_use_init_list = True

        # 生成代码
        code = self._generate_container_code(cpp_type, element_codes, can_use_init_list)
        return EmitResult(code=code, precedence=0)

    def _is_compatible_container_types(
        self, expected_type: str, actual_type: str
    ) -> bool:
        """检查两个容器类型是否兼容

        判断规则：
        1. 提取容器的基础模板名和元素类型
        2. 如果基础模板相同，且期望元素类型是 mc::variant，则兼容

        例如：
        - expected: mcpy::set<mc::variant>, actual: mcpy::set<int> -> True
        - expected: mc::array<mc::variant>, actual: mc::array<int> -> True
        - expected: mc::dict, actual: mcpy::set<int> -> False

        Args:
            expected_type: 期望的 C++ 类型
            actual_type: 实际的 C++ 类型

        Returns:
            如果类型兼容返回 True，否则返回 False
        """
        # 提取基础模板名
        expected_base = self._get_container_base_template(expected_type)
        actual_base = self._get_container_base_template(actual_type)

        if not expected_base or not actual_base or expected_base != actual_base:
            return False

        # 提取期望类型的元素类型
        expected_element = self._extract_element_type(expected_type)

        # 如果期望元素类型是 mc::variant，则任何元素类型都兼容
        return expected_element == "mc::variant"

    def _emit_tuple(
        self, node: Union[IRNode, IRTuple], context: CodeGenContext
    ) -> EmitResult:
        """生成元组代码

        根据目标类型生成不同的代码：
        - 强类型元组（std::tuple<Args...>）：使用 std::make_tuple
        - 弱类型元组（mc::variants）：使用初始化列表
        """

        generator = get_global_generator()

        # 类型断言
        tuple_node = node if isinstance(node, IRTuple) else None
        if tuple_node is None:
            return EmitResult(
                code=f"/* Invalid tuple node type: {type(node).__name__} */",
                precedence=0,
            )

        # 生成元素代码
        element_codes = []
        for element in tuple_node.elements:
            element_context = context.create_child_context(current_node=element)
            element_code = generator.generate_expression(element, element_context)
            element_codes.append(element_code)

        # 优先从类型推导获取类型（tuple 字面量的实际类型）
        inferred_type = node.get_extension_param("inferred_type")
        inferred_cpp_type = None
        if inferred_type and inferred_type.name == "tuple":
            # 有类型推导结果，使用类型系统获取 C++ 类型
            from mcpy.type_system import get_cpp_type

            inferred_cpp_type = get_cpp_type(inferred_type, node)

        # 获取目标类型（从赋值上下文传递）
        expected_cpp_type = context.expected_cpp_type
        if not expected_cpp_type:
            # 回退：从节点本身获取
            expected_cpp_type = context.get_cpp_type(node)

        # 确定最终使用的 C++ 类型
        # 如果推断类型是强类型 tuple，且目标类型是 mc::variant（Any），优先使用推断类型
        # 这样可以在赋值时进行类型转换，而不是直接生成弱类型
        # 但如果目标类型是 mc::variants（弱类型 tuple），则使用目标类型
        if (
            inferred_cpp_type
            and self._get_container_base_template(inferred_cpp_type) == "std::tuple"
            and expected_cpp_type == "mc::variant"
        ):
            # 推断类型是强类型 tuple，且目标类型是 mc::variant（Any），优先使用推断类型
            cpp_type = inferred_cpp_type
        elif expected_cpp_type:
            # 使用期望类型
            cpp_type = expected_cpp_type
        else:
            # 默认使用弱类型
            cpp_type = "mc::variants"

        # 获取容器基础模板名
        base_template = self._get_container_base_template(cpp_type)

        # 生成代码
        # tuple 比较特殊，需要区分强类型和弱类型
        # 强类型：std::tuple<Args...> -> 使用 std::make_tuple
        # 弱类型：mc::variants -> 使用初始化列表
        if base_template == "std::tuple":
            # 强类型元组：使用 std::make_tuple
            if element_codes:
                code = f"std::make_tuple({', '.join(element_codes)})"
            else:
                code = "std::tuple<>{}"
        elif cpp_type == "mc::variants" or (
            base_template and "variants" in base_template.lower()
        ):
            # 弱类型元组：使用 mc::variants 初始化列表
            if element_codes:
                code = self._generate_container_code(
                    "mc::variants", element_codes, can_use_init_list=False
                )
            else:
                code = "mc::variants{}"
        else:
            # 其他容器类型使用初始化列表（不支持省略类型）
            code = self._generate_container_code(
                cpp_type, element_codes, can_use_init_list=False
            )

        return EmitResult(code=code, precedence=0)

    def get_priority(self) -> int:
        """容器发射器的优先级"""
        return 6  # 低于操作符和内置函数，高于字面量
